import operator
import random
from typing import TypedDict

from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.types import Command
from typing_extensions import Annotated


class State(TypedDict):
    # NOTE: we define a reducer here
    foo: Annotated[str, operator.add]


def node_a(state: State):
    print("Called A")
    value = random.choice(["a", "b"])
    # this is a replacement for a conditional edge function
    if value == "a":
        goto = "node_b"
    else:
        goto = "node_c"

    # note how Command allows you to BOTH update the graph state AND route to the next node
    return Command(
        update={"foo": value},
        goto=goto,
        # this tells LangGraph to navigate to node_b or node_c in the parent graph
        # NOTE: this will navigate to the closest parent graph relative to the subgraph
        graph=Command.PARENT,
    )


def node_b(state: State):
    print("Called B")
    # NOTE: since we've defined a reducer, we don't need to manually append
    # new characters to existing 'foo' value. instead, reducer will append these
    # automatically (via operator.add)
    return {"foo": "b"}


def node_c(state: State):
    print("Called C")
    return {"foo": "c"}


if __name__ == '__main__':
    subgraph = StateGraph(State).add_node(node_a).add_edge(START, "node_a").compile()

    builder = StateGraph(State)
    builder.add_edge(START, "subgraph")
    builder.add_node("subgraph", subgraph)
    builder.add_node(node_b)
    builder.add_node(node_c)

    graph = builder.compile()
    result = graph.invoke({"foo": ""})
    print(result["foo"])

